#include "SpringSolver.h"
#include <set>
#include <map>
namespace common {

SpringSolver::SpringSolver(){

}

void SpringSolver::solve(){

	messages.clear();

	if(constraintDb==NULL || pointDb==NULL || primitiveDb==NULL || parameterDb==NULL){

		addMsg("Missing one of the DataBase");
		return;
	}


	_keyType* pointIdSet = pointDb->getIds();
	_keyType* constrainIdSet = constraintDb->getIds();
	_keyType* primitiveSet = primitiveDb->getIds();
	std::map<int,int> pointOrders = pointOrder(pointIdSet);

	int primSize = primitivesSize(primitiveSet);
	int conSize = constraintSize(constrainIdSet);

	ublasMatrix A(ublasMatrixZero(primSize + conSize, primSize + conSize));
	ublasMatrix invA(ublasMatrixZero(primSize + conSize, primSize + conSize));
	ublasVector dw(ublasVectorZero(primSize +conSize));
	ublasVector b(ublasVectorZero(primSize +conSize));
	ublasVector w(ublasVectorZero(primSize + conSize));

	ublasVectorRange dq(dw,range(0,primSize));
	ublasVectorRange da(dw,range(primSize,primSize + conSize));

	ublasVectorRange q(w,range(0,primSize));
	ublasVectorRange a(w,range(primSize,primSize + conSize));

	//prepare matrix for equation A*dw=b , we search for dw =  [dq,da]'
	// where dq - change on point position
	// and da - change on Lagrange'a multiplier

	project(A,range(0,primSize),range(0,primSize)) =*getMatrixK(primitiveSet) ;
	for(int i=0;i<5;i++){
		// fill matrix
		ublasMatrix JQ = *getJacobian(pointIdSet,constrainIdSet);
		ublasMatrix tJQ = trans(JQ);
		project(A,range(0,primSize),range(primSize,primSize+conSize)) = tJQ;
		project(A,range(primSize,primSize+conSize),range(0,primSize)) = JQ;

		project(b,range(0,primSize)) = getForceAll(primitiveSet);
		project(b,range(0,primSize)) -=prod(tJQ,a);
		project(b,range(primSize,primSize+conSize))= rhsQ(constrainIdSet)*-1;


		if(InvertMatrix(A,invA)){

			dw = prod(invA,b);
			w+=dw;
			updatePoints(dq,pointOrders);

		}else {

			std::cout << " Problem z inversja macierzy " << std::endl;
			break;
		}
	}


	skPvector(&q);
	skPvector(&a);

	delete pointIdSet ;
	delete constrainIdSet ;
	delete primitiveSet ;

}

ublasVector SpringSolver::rhsQ(_keyType* constrainIdSet){

	int sizeQ = constraintSize(constrainIdSet);
	ublasVector out(sizeQ);

	int sizeOfK =0;
	for (_keyType::iterator its = constrainIdSet->begin(); its != constrainIdSet->end(); its++){
		int sizeC = constraintDb->get(*its)->getSize();
		ublasVectorRange vr(out,range(sizeOfK,sizeOfK + sizeC));
		vr = constraintDb->get(*its)->getValue();
		sizeOfK += sizeC;
	}


	return out;

}

ublasVector SpringSolver::getForceAll(_keyType* primitiveSet) const {


	// what is the size of matrix
	int sizeOfK =primitivesSize(primitiveSet);

	ublasVector v(sizeOfK);


    //fill matrix with appropriate sub matrix for primitives
    int iK =0;
	for (_keyType::iterator its = primitiveSet->begin(); its != primitiveSet->end(); its++){

		int primSize = primitiveDb->get(*its)->getNumOfPoints()*2;

		project(v,range(iK,iK+primSize))=(primitiveDb->get(*its)->getForceAll());

		iK += primSize;
	}

	return v;

}

ublasMatrix* SpringSolver::getJacobian(_keyType* pointIdSet,_keyType* constrainIdSet) const{

	std::map<int,int> pointOrders =  pointOrder(pointIdSet);

	// matrix row size
	int sizeOfK =constraintSize(constrainIdSet);

	//prepare matrix -clear
	ublasMatrix* m = new ublasMatrix(ublasMatrixZero(sizeOfK,pointIdSet->size()*2));

    //fill matrix with appropriate sub matrix for constraints
    int i = 0; // row
	for (_keyType::iterator it = constrainIdSet->begin(); it != constrainIdSet->end(); it++) {
		int conSize = constraintDb->get(*it)->getSize();
		std::set<Point*> constrainPoints = constraintDb->get(*it)->getAssociatedPoint();
		for (std::set<Point*>::iterator iP = constrainPoints.begin(); iP != constrainPoints.end(); iP++) {
			//point position

			int j = pointOrders.find((*iP)->getId())->second;
			ublasMatrixRange mr((*m), range(i, i + conSize), range(j, j + 2));

			mr = constraintDb->get(*it)->getJacobianOnPoint(*iP);
		}
		i += conSize;
	}

    return m;
}

ublasMatrix* SpringSolver::getMatrixK(_keyType* primitiveSet) const {


	// what is the size of matrix
	int sizeOfK =primitivesSize(primitiveSet);

	//prepare matrix
	ublasMatrix* m = new ublasMatrix(ublasMatrixZero(sizeOfK,sizeOfK));

    //fill matrix with appropriate sub matrix for primitives
    int iK =0;
	for (_keyType::iterator its = primitiveSet->begin(); its != primitiveSet->end(); its++){

		int primSize = primitiveDb->get(*its)->getNumOfPoints()*2;

		ublasMatrixRange mr ((*m), range (iK,iK + primSize), range (iK, iK + primSize));

		mr=(primitiveDb->get(*its)->getJacobianAll()); //FIXME look out on types

		iK += primSize;
	}

	return m;
}

int SpringSolver::constraintSize(_keyType* constrainIdSet) const{

	// matrix row size
	int sizeOfK =0;
	for (_keyType::iterator its = constrainIdSet->begin(); its != constrainIdSet->end(); its++){
		sizeOfK += constraintDb->get(*its)->getSize();
	}
	return sizeOfK;
}

int SpringSolver::primitivesSize(_keyType* primitiveSet) const{
	int sizeOfK =0;
	for (_keyType::iterator its = primitiveSet->begin(); its != primitiveSet->end(); its++){
		sizeOfK += primitiveDb->get(*its)->getNumOfPoints()*2;
	}
	return sizeOfK;
}

std::map<int,int> SpringSolver::pointOrder(_keyType* pointIdSet) const{

	std::map<int,int> pointOrder; // first: pointId , second: position
	int poz =0;
	for(_keyType::iterator io=pointIdSet->begin(); io!= pointIdSet->end(); io++){
		pointOrder.insert(make_pair(*io,poz));
		poz+=2;
	}
	return pointOrder;
}
void SpringSolver::updatePoints(ublasVectorRange& dq,std::map<int,int>& pointOrders){

	for(std::map<int,int>::iterator it = pointOrders.begin(); it!= pointOrders.end(); it++){
		Point* pn = pointDb->get(it->first);
		//FIXME PLUS or MINUS
		pn->setX(dq(it->second) + pn->getX());
		pn->setY(dq(it->second+1) + pn->getY());
	}

}
void SpringSolver::relax(){

	messages.clear();

	//reset all anchors

	_keyType* ids = pointDb->getIds();
	for(_keyType::iterator it =ids->begin();it!=ids->end();it++){
		pointDb->get(*it)->resetAnchor();
	}

	addMsg("Points relaxed");

}

SpringSolver::~SpringSolver() {
}

}
